###競品ner識別模型

import os
import pickle
import re
from typing import List


def get_root_dir(path):
    path_list=path.split(os.path.sep)
    index=path_list.index("featurelib")
    return os.path.sep.join(path_list[:index+1])


class SmsEsLoan0V1NerModel:
    def __init__(self, country_type="es"):
        self.ROOT_DIR = get_root_dir(os.path.abspath("."))
        self.CONF_DIR = os.path.join(
            self.ROOT_DIR, "feature_conf", "sms", "un", "sms_es_loan0_v1"
        )
        self.MODEL_CONF_DIR = os.path.join(self.CONF_DIR, "model_conf")
        pkl_name = "spanish_cmp_ner.pkl" if country_type == "es" else "mx_cmp_ner.pkl"
        with open(os.path.join(self.MODEL_CONF_DIR, pkl_name), "rb") as f:
            self.nlp = pickle.load(f)

    def predict(self, inputs: List[str]):
        all_entities = []
        for msg in inputs:
            doc = self.nlp(msg)
            entities = {}
            for ent in doc.ents:
                entities[ent.label_] = ent.text
            all_entities.append(entities)
        return all_entities
